from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
import os
import tempfile
import logging
import uvicorn

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="语音转文本 API")

# 初始化模型
model_dir = "iic/SenseVoiceSmall"
try:
    model = AutoModel(
        model=model_dir,
        vad_model="fsmn-vad",
        vad_kwargs={"max_single_segment_time": 30000},
        device="cpu",
        disable_update=True,  # 禁用更新检查，加快启动
    )
    logger.info("模型初始化成功")
except Exception as e:
    logger.error(f"模型初始化失败: {str(e)}")
    raise e

@app.post("/transcribe")
async def transcribe_audio(file: UploadFile = File(...)):
    try:
        # 验证文件类型
        if not file.filename.endswith(('.wav', '.mp3', '.flac')):
            return JSONResponse(
                status_code=400,
                content={"error": "只支持 WAV、MP3 或 FLAC 格式的音频文件"}
            )

        # 创建临时文件保存上传的音频
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
            temp_file.write(await file.read())
            temp_file_path = temp_file.name

        try:
            # 使用模型进行转录
            res = model.generate(
                input=temp_file_path,
                cache={},
                language="auto",
                use_itn=True,
                batch_size_s=60,
                merge_vad=True,
                merge_length_s=15,
            )

            # 处理转录结果
            text = rich_transcription_postprocess(res[0]["text"])
            logger.info(f"转录成功: {text}")

            return JSONResponse(
                status_code=200,
                content={"text": text}
            )

        finally:
            # 清理临时文件
            if os.path.exists(temp_file_path):
                os.unlink(temp_file_path)

    except Exception as e:
        logger.error(f"转录过程中发生错误: {str(e)}")
        return JSONResponse(
            status_code=500,
            content={"error": f"转录失败: {str(e)}"}
        )

@app.get("/health")
async def health_check():
    return JSONResponse(
        status_code=200,
        content={"status": "healthy"}
    )

# 添加启动代码
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)